
package org.apache.solr.servlet;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.io.Writer;
import java.nio.charset.Charset;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.WeakHashMap;
import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.xml.parsers.ParserConfigurationException;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.cloud.ClusterState;
import org.apache.solr.common.cloud.Replica;
import org.apache.solr.common.cloud.Slice;
import org.apache.solr.common.cloud.ZkNodeProps;
import org.apache.solr.common.cloud.ZkStateReader;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.util.ContentStreamBase;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.common.util.SimpleOrderedMap;
import org.apache.solr.core.Config;
import org.apache.solr.core.CoreContainer;
import org.apache.solr.core.SolrConfig;
import org.apache.solr.core.SolrCore;
import org.apache.solr.handler.ContentStreamHandlerBase;
import org.apache.solr.request.ServletSolrParams;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.request.SolrQueryRequestBase;
import org.apache.solr.request.SolrRequestHandler;
import org.apache.solr.request.SolrRequestInfo;
import org.apache.solr.response.BinaryQueryResponseWriter;
import org.apache.solr.response.QueryResponseWriter;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.servlet.cache.HttpCacheHeaderUtil;
import org.apache.solr.servlet.cache.Method;
import org.apache.solr.util.FastWriter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

/**
 * This filter looks at the incoming URL maps them to handlers defined in solrconfig.xml
 *
 * @since solr 1.2
 */
public class SolrDispatchFilter implements Filter {

    final Logger log = LoggerFactory.getLogger(SolrDispatchFilter.class);
    protected volatile CoreContainer cores;
    protected String pathPrefix = null; // strip this from the beginning of a path
    protected String abortErrorMessage = null;
    protected final Map<SolrConfig, SolrRequestParsers> parsers = new WeakHashMap<>();
    protected final SolrRequestParsers adminRequestParser;
    private static final Charset UTF8 = Charset.forName("UTF-8");

    public SolrDispatchFilter() {

        try {
            adminRequestParser = new SolrRequestParsers(new Config(null, "solr", new InputSource(new ByteArrayInputStream("<root/>".getBytes("UTF-8"))), ""));
        }
        catch (ParserConfigurationException | IOException | SAXException e) {
            //unlikely
            throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, e);
        }
    }

    @Override
    public void init(FilterConfig config) throws ServletException {

        log.info("SolrDispatchFilter.init()");

        CoreContainer.Initializer init = createInitializer();
        try {
            // web.xml configuration
            this.pathPrefix = config.getInitParameter("path-prefix");

            this.cores = init.initialize();
            log.info("user.dir=" + System.getProperty("user.dir"));
        }
        catch (IOException | ParserConfigurationException | SAXException t) {
            // catch this so our filter still works
            log.error("Could not start Solr. Check solr/home property and the logs");
            SolrCore.log(t);
        }

        log.info("SolrDispatchFilter.init() done");
    }

    public CoreContainer getCores() {
        return cores;
    }

    /**
     * Method to override to change how CoreContainer initialization is
     * performed.
     */
    protected CoreContainer.Initializer createInitializer() {
        return new CoreContainer.Initializer();
    }

    @Override
    public void destroy() {

        if (cores != null) {
            cores.shutdown();
            cores = null;
        }
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {

        if (abortErrorMessage != null) {
            ((HttpServletResponse) response).sendError(500, abortErrorMessage);
            return;
        }

        if (this.cores == null) {
            ((HttpServletResponse) response).sendError(503, "Server is shutting down");
            return;
        }

        CoreContainer lcores = this.cores;
        SolrCore core = null;
        SolrQueryRequest solrReq = null;

        if (request instanceof HttpServletRequest) {
            HttpServletRequest req = (HttpServletRequest) request;
            HttpServletResponse resp = (HttpServletResponse) response;
            SolrRequestHandler handler = null;
            String corename = "";
            try {
                // put the core container in request attribute
                req.setAttribute("org.apache.solr.CoreContainer", lcores);
                String path = req.getServletPath();
                if (req.getPathInfo() != null) {
                    // this lets you handle /update/commit when /update is a servlet
                    path += req.getPathInfo();
                }
                if (pathPrefix != null && path.startsWith(pathPrefix)) {
                    path = path.substring(pathPrefix.length());
                }
                // check for management path
                String alternate = lcores.getManagementPath();
                if (alternate != null && path.startsWith(alternate)) {
                    path = path.substring(0, alternate.length());
                }
                // unused feature ?
                int idx = path.indexOf(':');
                if (idx > 0) {
                    // save the portion after the ':' for a 'handler' path parameter
                    path = path.substring(0, idx);
                }

                // Check for the core admin page
                if (path.equals(lcores.getAdminPath())) {
                    handler = lcores.getMultiCoreHandler();
                    solrReq = adminRequestParser.parse(null, path, req);
                    handleAdminRequest(req, response, handler, solrReq);
                    return;
                }
                // Check for the core admin collections url
                if (path.equals("/admin/collections")) {
                    handler = lcores.getCollectionsHandler();
                    solrReq = adminRequestParser.parse(null, path, req);
                    handleAdminRequest(req, response, handler, solrReq);
                    return;
                }
                else {
                    //otherwise, we should find a core from the path
                    idx = path.indexOf("/", 1);
                    if (idx > 1) {
                        // try to get the corename as a request parameter first
                        corename = path.substring(1, idx);
                        core = lcores.getCore(corename);
                        if (core != null) {
                            path = path.substring(idx);
                        }
                    }
                    if (core == null) {
                        if (!lcores.isZooKeeperAware()) {
                            core = lcores.getCore("");
                        }
                    }
                }

                if (core == null && lcores.isZooKeeperAware()) {
                    // we couldn't find the core - lets make sure a collection was not specified instead
                    core = getCoreByCollection(lcores, corename, path);

                    if (core != null) {
                        // we found a core, update the path
                        path = path.substring(idx);
                    }
                    else {
                        // try the default core
                        core = lcores.getCore("");
                    }
                    // TODO: if we couldn't find it locally, look on other nodes
                }

                // With a valid core...
                if (core != null) {
                    final SolrConfig config = core.getSolrConfig();
                    // get or create/cache the parser for the core
                    SolrRequestParsers parser = parsers.get(config);
                    if (parser == null) {
                        parser = new SolrRequestParsers(config);
                        parsers.put(config, parser);
                    }

                    // Determine the handler from the url path if not set
                    // (we might already have selected the cores handler)
                    if (handler == null && path.length() > 1) { // don't match "" or "/" as valid path
                        handler = core.getRequestHandler(path);
                        // no handler yet but allowed to handle select; let's check
                        if (handler == null && parser.isHandleSelect()) {
                            if ("/select".equals(path) || "/select/".equals(path)) {
                                solrReq = parser.parse(core, path, req);
                                String qt = solrReq.getParams().get(CommonParams.QT);
                                handler = core.getRequestHandler(qt);
                                if (handler == null) {
                                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "unknown handler: " + qt);
                                }
                                if (qt != null && qt.startsWith("/") && (handler instanceof ContentStreamHandlerBase)) {
                                    //For security reasons it's a bad idea to allow a leading '/', ex: /select?qt=/update see SOLR-3161
                                    //There was no restriction from Solr 1.4 thru 3.5 and it's not supported for update handlers.
                                    throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid Request Handler ('qt').  Do not use /select to access: " + qt);
                                }
                            }
                        }
                    }

                    // With a valid handler and a valid core...
                    if (handler != null) {
                        // if not a /select, create the request
                        if (solrReq == null) {
                            solrReq = parser.parse(core, path, req);
                        }

                        final Method reqMethod = Method.getMethod(req.getMethod());
                        HttpCacheHeaderUtil.setCacheControlHeader(config, resp, reqMethod);
                        // unless we have been explicitly told not to, do cache validation
                        // if we fail cache validation, execute the query
                        if (config.getHttpCachingConfig().isNever304() || !HttpCacheHeaderUtil.doCacheHeaderValidation(solrReq, req, reqMethod, resp)) {
                            SolrQueryResponse solrRsp = new SolrQueryResponse();
                            /* even for HEAD requests, we need to execute the handler to
                             * ensure we don't get an error (and to make sure the correct
                             * QueryResponseWriter is selected and we get the correct
                             * Content-Type)
                             */
                            SolrRequestInfo.setRequestInfo(new SolrRequestInfo(solrReq, solrRsp));
                            this.execute(req, handler, solrReq, solrRsp);
                            HttpCacheHeaderUtil.checkHttpCachingVeto(solrRsp, resp, reqMethod);
                            // add info to http headers
                            //TODO: See SOLR-232 and SOLR-267.  
                            /*try {
                             NamedList solrRspHeader = solrRsp.getResponseHeader();
                             for (int i=0; i<solrRspHeader.size(); i++) {
                             ((javax.servlet.http.HttpServletResponse) response).addHeader(("Solr-" + solrRspHeader.getName(i)), String.valueOf(solrRspHeader.getVal(i)));
                             }
                             } catch (ClassCastException cce) {
                             log.log(Level.WARNING, "exception adding response header log information", cce);
                             }*/
                            QueryResponseWriter responseWriter = core.getQueryResponseWriter(solrReq);
                            writeResponse(solrRsp, response, responseWriter, solrReq, reqMethod);
                        }
                        return; // we are done with a valid handler
                    }
                }
                log.debug("no handler or core retrieved for " + path + ", follow through...");
            }
            catch (Throwable ex) {
                sendError(core, solrReq, request, (HttpServletResponse) response, ex);
                return;
            }
            finally {
                if (solrReq != null) {
                    solrReq.close();
                }
                if (core != null) {
                    core.close();
                }
                SolrRequestInfo.clearRequestInfo();
            }
        }

        // Otherwise let the webapp handle the request
        chain.doFilter(request, response);
    }

    private SolrCore getCoreByCollection(CoreContainer cores, String corename, String path) {

        String collection = corename;
        ZkStateReader zkStateReader = cores.getZkController().getZkStateReader();

        ClusterState clusterState = zkStateReader.getClusterState();
        Map<String, Slice> slices = clusterState.getSlices(collection);
        if (slices == null) {
            return null;
        }
        // look for a core on this node
        Set<Entry<String, Slice>> entries = slices.entrySet();
        SolrCore core = null;
        done:
        for (Entry<String, Slice> entry : entries) {
            // first see if we have the leader
            ZkNodeProps leaderProps = clusterState.getLeader(collection, entry.getKey());
            if (leaderProps != null) {
                core = checkProps(cores, path, leaderProps);
            }
            if (core != null) {
                break done;
            }

            // check everyone then
            Map<String, Replica> shards = entry.getValue().getReplicasMap();
            Set<Entry<String, Replica>> shardEntries = shards.entrySet();
            for (Entry<String, Replica> shardEntry : shardEntries) {
                Replica zkProps = shardEntry.getValue();
                core = checkProps(cores, path, zkProps);
                if (core != null) {
                    break done;
                }
            }
        }
        return core;
    }

    private SolrCore checkProps(CoreContainer cores, String path, ZkNodeProps zkProps) {

        String corename;
        SolrCore core = null;
        if (cores.getZkController().getNodeName().equals(zkProps.getStr(ZkStateReader.NODE_NAME_PROP))) {
            corename = zkProps.getStr(ZkStateReader.CORE_NAME_PROP);
            core = cores.getCore(corename);
        }
        return core;
    }

    private void handleAdminRequest(HttpServletRequest req, ServletResponse response, SolrRequestHandler handler, SolrQueryRequest solrReq) throws IOException {

        SolrQueryResponse solrResp = new SolrQueryResponse();
        final NamedList<Object> responseHeader = new SimpleOrderedMap<>();
        solrResp.add("responseHeader", responseHeader);
        NamedList toLog = solrResp.getToLog();
        toLog.add("webapp", req.getContextPath());
        toLog.add("path", solrReq.getContext().get("path"));
        toLog.add("params", "{" + solrReq.getParamString() + "}");
        handler.handleRequest(solrReq, solrResp);
        SolrCore.setResponseHeaderValues(handler, solrReq, solrResp);
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < toLog.size(); i++) {
            String name = toLog.getName(i);
            Object val = toLog.getVal(i);
            sb.append(name).append("=").append(val).append(" ");
        }
        QueryResponseWriter respWriter = SolrCore.DEFAULT_RESPONSE_WRITERS.get(solrReq.getParams().get(CommonParams.WT));
        if (respWriter == null) {
            respWriter = SolrCore.DEFAULT_RESPONSE_WRITERS.get("standard");
        }
        writeResponse(solrResp, response, respWriter, solrReq, Method.getMethod(req.getMethod()));
    }

    private void writeResponse(SolrQueryResponse solrRsp, ServletResponse response, QueryResponseWriter responseWriter, SolrQueryRequest solrReq, Method reqMethod) throws IOException {

        // Now write it out
        final String ct = responseWriter.getContentType(solrReq, solrRsp);
        // don't call setContentType on null
        if (null != ct) {
            response.setContentType(ct);
        }

        if (solrRsp.getException() != null) {
            NamedList info = new SimpleOrderedMap();
            int code = getErrorInfo(solrRsp.getException(), info);
            solrRsp.add("error", info);
            ((HttpServletResponse) response).setStatus(code);
        }

        if (Method.HEAD != reqMethod) {
            if (responseWriter instanceof BinaryQueryResponseWriter) {
                BinaryQueryResponseWriter binWriter = (BinaryQueryResponseWriter) responseWriter;
                binWriter.write(response.getOutputStream(), solrReq, solrRsp);
            }
            else {
                String charset = ContentStreamBase.getCharsetFromContentType(ct);
                Writer out = (charset == null || charset.equalsIgnoreCase("UTF-8"))
                        ? new OutputStreamWriter(response.getOutputStream(), UTF8)
                        : new OutputStreamWriter(response.getOutputStream(), charset);
                out = new FastWriter(out);
                responseWriter.write(out, solrReq, solrRsp);
                out.flush();
            }
        }
        //else http HEAD request, nothing to write out, waited this long just to get ContentType
    }

    protected int getErrorInfo(Throwable ex, NamedList info) {

        int code = 500;
        if (ex instanceof SolrException) {
            code = ((SolrException) ex).code();
        }

        String msg = null;
        for (Throwable th = ex; th != null; th = th.getCause()) {
            msg = th.getMessage();
            if (msg != null) {
                break;
            }
        }
        if (msg != null) {
            info.add("msg", msg);
        }

        // For any regular code, don't include the stack trace
        if (code == 500 || code < 100) {
            StringWriter sw = new StringWriter();
            ex.printStackTrace(new PrintWriter(sw));
            SolrException.log(log, null, ex);
            info.add("trace", sw.toString());

            // non standard codes have undefined results with various servers
            if (code < 100) {
                log.warn("invalid return code: " + code);
                code = 500;
            }
        }
        info.add("code", new Integer(code));
        return code;
    }

    protected void execute(HttpServletRequest req, SolrRequestHandler handler, SolrQueryRequest sreq, SolrQueryResponse rsp) {
        // a custom filter could add more stuff to the request before passing it on.
        // for example: sreq.getContext().put( "HttpServletRequest", req );
        // used for logging query stats in SolrCore.execute()
        sreq.getContext().put("webapp", req.getContextPath());
        sreq.getCore().execute(handler, sreq, rsp);
    }

    protected void sendError(SolrCore core, SolrQueryRequest req, ServletRequest request, HttpServletResponse response, Throwable ex) throws IOException {

        try {
            SolrQueryResponse solrResp = new SolrQueryResponse();
            if (ex instanceof Exception) {
                solrResp.setException((Exception) ex);
            }
            else {
                solrResp.setException(new RuntimeException(ex));
            }
            if (core == null) {
                core = cores.getCore(""); // default core
            }
            if (req == null) {
                req = new SolrQueryRequestBase(core, new ServletSolrParams(request)) {
                };
            }
            QueryResponseWriter writer = core.getQueryResponseWriter(req);
            writeResponse(solrResp, response, writer, req, Method.GET);
        }
        catch (Throwable t) { // This error really does not matter
            SimpleOrderedMap info = new SimpleOrderedMap();
            int code = getErrorInfo(ex, info);
            response.sendError(code, info.toString());
        }
    }

    /**
     * Set the prefix for all paths. This is useful if you want to apply the
     * filter to something other then /*, perhaps because you are merging this
     * filter into a larger web application.
     *
     * For example, if web.xml specifies:
     *
     * <filter-mapping> <filter-name>SolrRequestFilter</filter-name>
     * <url-pattern>/xxx/*</url-pattern> </filter-mapping>
     *
     * Make sure to set the PathPrefix to "/xxx" either with this function or in
     * web.xml.
     *
     * <init-param> <param-name>path-prefix</param-name>
     * <param-value>/xxx</param-value> </init-param>
     *
     */
    public void setPathPrefix(String pathPrefix) {
        this.pathPrefix = pathPrefix;
    }

    public String getPathPrefix() {
        return pathPrefix;
    }
}
